import pandas
import os
import matplotlib.pyplot as plt
import numpy as np

# 读取 CSV 数据
baseDir = './2025_Problem_C_Data'
athletes = pandas.read_csv(os.path.join(baseDir, 'summerOly_athletes.csv'), encoding='utf-8', low_memory=False)
hosts = pandas.read_csv(os.path.join(baseDir, 'summerOly_hosts.csv'), encoding='utf-8', low_memory=False)
medals = pandas.read_csv(os.path.join(baseDir, 'summerOly_medal_counts.csv'), encoding='utf-8', low_memory=False)
programs = pandas.read_csv(os.path.join(baseDir, 'summerOly_programs.csv'), encoding='windows 1256', low_memory=False)

## 数据预处理
medals['NOC'] = medals['NOC'].str.strip()  # medals 去除 NOC 字段前后的空格
medals = medals.groupby('NOC').filter(lambda x: x['Year'].max() >= 2024)  # medals 中筛掉 2024 没有参赛的队伍

# medals 按照 NOC 字段分组
medals_grouped = medals.groupby('NOC')
print(medals_grouped.size())
# athletes 按照 NOC 字段分组
athletes_grouped = athletes.groupby('NOC')
print(athletes_grouped.size())

# 找到 2024 年之后参赛 5 次以上的队伍的数据
times = 5
medals_class_2 = medals[medals['Year'] > 2024 - times * 4].groupby('NOC').filter(lambda x: len(x) >= times)
# 计算这几个队伍的平均奖牌数和标准差
statistic = medals_class_2[['NOC', 'Rank', 'Gold', 'Silver', 'Bronze', 'Total']].groupby('NOC').agg(['mean', 'std'])
print(statistic)
